#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy, re, csv
from pyscf import gto, lib, scf, symm, ao2mo
from pyscf_ccso import mo_sym, ri_blocks
from scipy.io import FortranFile

def putallso(nirrep, ilist, filename, iin):
   isym, numl, numr, nsize = mo_sym.geteridata(nirrep, ilist, filename)
   print('numl and numr', numl, numr)
   ioff = 0
   for irrep in range(nirrep):
       itmp = mo_sym.putlstso(nirrep, 1, numr[irrep], numl[irrep], irrep+1, 1, filename, 
                         iin[ioff:ioff+numl[irrep]*numr[irrep]])
       ioff = ioff + numl[irrep]*numr[irrep]

def getallso(nirrep, ilist, filename):
   isym, numl, numr, nsize = mo_sym.geteridata(nirrep, ilist, filename)
   iout = []
   for irrep in range(nirrep):
       itmp = mo_sym.getlstso(nirrep, 1, numr[irrep], numl[irrep], irrep+1, 1, filename)
       iout.extend(itmp)
   return iout

def symminfo(groupname):
   if groupname == 'Dooh' or groupname == 'SO3':
      group_corr = 'D2h'
   elif groupname == 'Coov':
      group_corr = 'C2v'
   else:
      group_corr = groupname
   symm_table = symm.param.IRREP_ID_TABLE[group_corr]
#   print('symm_table', symm_table)
   nirrep = len(symm_table)
   irrep_sym = ['']*nirrep
   for irsymb in symm_table:
      irid = symm_table[irsymb]
      irrep_sym[irid] = irsymb
   return group_corr, nirrep, irrep_sym

def symrxyz(group_corr):
   group1 = ['D2h', 'C2v']
   group2 = ['C2h', 'Cs', 'C2']
   group3 = ['Ci', 'C1']
   group4 = ['D2']
   if group_corr in group1:
      symrx = 2; symry = 3; symrz = 1
   elif group_corr in group2:
      symrx = 1; symry = 1; symrz = 0
   elif group_corr in group3:
      symrx = 0; symry = 0; symrz = 0
   elif group_corr in group4:
      symrx = 3; symry = 2; symrz = 1
   return symrx, symry, symrz

def getioffset(nirrep, occno):
   ioffset = numpy.zeros(nirrep+1, dtype = int)
   for irrep in range(nirrep):
      ioffset[irrep+1] = ioffset[irrep] + occno[irrep]
   return ioffset

def symorb_sort0(mol, nirrep, mo_occ, mo_coeff, symid, frozen=0):
   occidx = mo_occ[frozen:] == 2
   viridx = mo_occ[frozen:] == 0
   nmo = len(mo_occ) - frozen
   nocc = numpy.count_nonzero(occidx)
   nvir = numpy.count_nonzero(viridx)
   isymmo  = numpy.zeros((nmo,8), dtype = int)
   pop = numpy.zeros(8, dtype = int)
   vrt = numpy.zeros(8, dtype = int)
   for irrep in range(nirrep):
       idx0 = symid == irrep
       idx1 = occidx & idx0; idx2 = viridx & idx0
       irrep_occ = list(numpy.arange(nmo)[idx1])
       irrep_vir = list(numpy.arange(nmo)[idx2])
       irrep_mo  = irrep_occ + irrep_vir
       pop[irrep] = len(irrep_occ)
       vrt[irrep] = len(irrep_vir)
       isymmo[0:pop[irrep]+vrt[irrep], irrep] = irrep_mo
#  print(symname) 
#  print(symid2) 
   return pop, vrt, isymmo

def symorb_sort(mol, nirrep, mo_occ, mo_coeff, frozen=0):
#   print(mo_occ)
   occidx = mo_occ[frozen:] == 2
   viridx = mo_occ[frozen:] == 0
#   print(occidx, viridx)
   nmo = len(mo_occ) - frozen
   nocc = numpy.count_nonzero(occidx)
   nvir = numpy.count_nonzero(viridx)
   isymmo  = numpy.zeros((nmo,8), dtype = int)
   pop = numpy.zeros(8, dtype = int)
   vrt = numpy.zeros(8, dtype = int)
   if nirrep==1:
      pop[0]=nocc
      vrt[0]=nvir
      isymmo[:, 0] = numpy.arange(nmo)
   elif nirrep>1:
      symid2 = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mo_coeff[:,frozen:])
      symid = symid2 % 10
      for irrep in range(nirrep):
          idx0 = symid == irrep
          idx1 = occidx & idx0; idx2 = viridx & idx0
#         print('idx0', idx0)
#         print('idx1', idx1)
#         print('idx2', idx2)
          irrep_occ = list(numpy.arange(nmo)[idx1])
#         print(irrep_occ)
          irrep_vir = list(numpy.arange(nmo)[idx2])
          irrep_mo  = irrep_occ + irrep_vir
          pop[irrep] = len(irrep_occ)
          vrt[irrep] = len(irrep_vir)
          isymmo[0:pop[irrep]+vrt[irrep], irrep] = irrep_mo
   return pop, vrt, isymmo

def symorb_sort_u(mol, nirrep, mo_occ, mo_coeff, frozen=0):
#   print(mo_occ)
   occidx_a = mo_occ[0][frozen:] == 1
   viridx_a = mo_occ[0][frozen:] == 0
   occidx_b = mo_occ[1][frozen:] == 1
   viridx_b = mo_occ[1][frozen:] == 0
#   print(occidx_a, occidx_b)
#   print(viridx_a, viridx_b)
   nmo = len(mo_occ[0]) - frozen
   nocc_a = numpy.count_nonzero(occidx_a)
   nvir_a = numpy.count_nonzero(viridx_a)
   nocc_b = numpy.count_nonzero(occidx_b)
   nvir_b = numpy.count_nonzero(viridx_b)
#   isymmo_a  = numpy.zeros((nmo,8), dtype = int)
#   isymmo_b  = numpy.zeros((nmo,8), dtype = int)
   isymmo  = numpy.zeros((nmo,8,2), dtype = int)
   pop = numpy.zeros((8,2), dtype = int)
   vrt = numpy.zeros((8,2), dtype = int)

   if nirrep==1:
      pop[0,0]=nocc_a
      vrt[0,0]=nvir_a
      pop[0,1]=nocc_b
      vrt[0,1]=nvir_b
      isymmo[:,0,0] = numpy.arange(nmo)
      isymmo[:,0,1] = numpy.arange(nmo)
   elif nirrep>1:
      symid_a = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mo_coeff[0][:,frozen:])
      symid_b = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mo_coeff[1][:,frozen:])
#      print(symid_a)
#      print(symid_b)
      for irrep in range(nirrep):
#=======define alpha spin
          idx0 = symid_a == irrep
          idx1 = occidx_a & idx0; idx2 = viridx_a & idx0
          irrep_occ = list(numpy.arange(nmo)[idx1])
          irrep_vir = list(numpy.arange(nmo)[idx2])
          irrep_mo  = irrep_occ + irrep_vir
          pop[irrep,0] = len(irrep_occ)
          vrt[irrep,0] = len(irrep_vir)
          isymmo[0:pop[irrep,0]+vrt[irrep,0], irrep,0] = irrep_mo
#          isymmo_a[0:pop[irrep,0]+vrt[irrep,0], irrep] = irrep_mo
#=======define beta spin
          idx0 = symid_b == irrep
          idx1 = occidx_b & idx0; idx2 = viridx_b & idx0
          irrep_occ = list(numpy.arange(nmo)[idx1])
          irrep_vir = list(numpy.arange(nmo)[idx2])
          irrep_mo  = irrep_occ + irrep_vir
          pop[irrep,1] = len(irrep_occ)
          vrt[irrep,1] = len(irrep_vir)
          isymmo[0:pop[irrep,0]+vrt[irrep,0], irrep,1] = irrep_mo
#          isymmo_b[0:pop[irrep,0]+vrt[irrep,0], irrep] = irrep_mo
#      print('alpha', isymmo_a)
#      print('beta ', isymmo_b)

   return pop, vrt, isymmo
#   exit() 

def symorb_sort2(mol, nirrep, mo_occ, mo_coeff, frozen):
   listocca = list(numpy.argwhere(mo_occ > 0).flatten())
   listvira = list(numpy.argwhere(mo_occ == 0).flatten())
#   print('listocca', listocca)
   listocca.sort()
   listvira.sort()
#   print('listocca', listocca)
   listorba = listocca + listvira
   nocca = len(listocca)-frozen
   nvira = len(listvira)
   nmo = nocca+nvira
#  print('nocca, nvira', nocca, nvira)
#  print('listorba', listorba)
#  print('shape of mo_coeff', mo_coeff.shape)
   mocoeffa = mo_coeff[:,listorba]
   symnamea = symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mocoeffa)
   symida = symm.label_orb_symm(mol, mol.irrep_id,     mol.symm_orb, mocoeffa)
#  symnamea = symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mocoeffa, check=False)
#  symida = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mocoeffa, check=False)
#   print('irrep name:', symnamea)
#   print('irrep   id:', symida)
   symida = symida % 10
#  print('old irrep id:', symida)
   symid = symida[frozen:]
#  for i in range(nmo):
#      symid[i] = symid[i]-frozen
   symida = symid
#  print('new irrep id:', symida)
#  print('leth:', len(symida), nmo)
   sort_occ_a = []
   sort_vir_a = []
   sort_all_a = []
   pop = numpy.zeros(8, dtype = int)
   vrt = numpy.zeros(8, dtype = int)
   for irrep in range(nirrep):
       irrep_occa = list(numpy.argwhere(symida[:nocca] == irrep).flatten())
       irrep_vira = list(numpy.argwhere(symida[nocca:] == irrep).flatten() + nocca)
       irrep_occa.sort()
       irrep_vira.sort()
       pop[irrep] = len(irrep_occa)
       vrt[irrep] = len(irrep_vira)
       sort_occ_a += irrep_occa
       sort_vir_a += irrep_vira
       sort_all_a += irrep_occa + irrep_vira
#      print('irrep, irrep_occa', irrep, irrep_occa)
#      print('irrep, irrep_vira', irrep, irrep_vira)
#  print('sort_occ_a', sort_occ_a)
#  print('sort_vir_a', sort_vir_a)
   return pop, vrt, sort_occ_a, sort_vir_a, sort_all_a, mocoeffa

def checksym(matin, nirrep, dirprd, ioffidxi, ioffidxj, symr, symstr):
   symidx = []
   for irrepi in range(nirrep):
      i1 = ioffidxi[irrepi]; i2 = ioffidxi[irrepi+1]
      if i1 == i2:
         continue
      for irrepj in range(nirrep):
         j1 = ioffidxj[irrepj]; j2 = ioffidxj[irrepj+1]
         if j1 == j2:
            continue
#         print('irrepj', i1, i2, j1, j2, irrepj, nirrep)
         if numpy.amax(matin[i1:i2,j1:j2])>1.e-8:
            symin = dirprd[irrepi, irrepj]
            if symin not in symidx:
               symidx.append(symin)
   if len(symidx)>1:
      print('INCORRECT Symmetry! more than two symmetry for ', symstr)
   elif len(symidx)==1 and symidx[0] != symr:
      print('INCORRECT Symmetry! ', symstr, symr, symidx[0])

def checksym2(matin, nirrep, dirprd, ioffidxi, ioffidxj, symr, symstr):
   symidx = []
   err_sym = 0
   for irrepi in range(nirrep):
      i1 = ioffidxi[irrepi]; i2 = ioffidxi[irrepi+1]
      if i1 == i2:
         continue
      for irrepj in range(nirrep):
         j1 = ioffidxj[irrepj]; j2 = ioffidxj[irrepj+1]
         if j1 == j2:
            continue
#         print('irrepj', i1, i2, j1, j2, irrepj, nirrep)
         if numpy.amax(matin[i1:i2,j1:j2])>1.e-8:
            symin = dirprd[irrepi, irrepj]
            if symin not in symidx:
               symidx.append(symin)
   if len(symidx)>1:
      err_sym = 1
      print('INCORRECT Symmetry! more than two symmetry for ', symstr)
   elif len(symidx)==1 and symidx[0] != symr:
      err_sym = 1
      print('INCORRECT Symmetry! ', symstr, symr, symidx[0])
   return err_sym

def readinp():
   with open('CCSOINP', 'r') as ccsoinp:
       lists = ccsoinp.readlines()
   lth = len(lists)
   iop_nbo_state = 0
   iop_nbo_act = 1
   iop_soc = 0 
   iop_ri = 0
   iop_nbo_state = 0
   iop_nbo_act = 0
   thresh = 0.0
   for i in range(0,lth):
   
#    print(lists[i])
#
       if lists[i][0]=='!' or lists[i][0]=='#':
#           print('continue')
           continue
       if 'nbo_state' in lists[i]:
           if 'ground' in lists[i]:
               iop_nbo_state = 1
           elif 'excited' in lists[i]:
               iop_nbo_state = 2
           elif 'none' in lists[i]:
               iop_nbo_state = 0

       if 'nbo_act' in lists[i]:
           if 'single' in lists[i]:
               iop_nbo_act = 0
           elif 'total' in lists[i]:
               iop_nbo_act = 1

       if 'ri=' in lists[i]:
           if 'on' in lists[i]:
               iop_ri = 1  
           elif 'off' in lists[i]: 
               iop_ri = 0 
        
       if 'pysoc' in lists[i]:
           if 'on' in lists[i]:
              iop_soc = 1
       if 'nbo_value' in lists[i]:
           tmp_value = str(lists[i])
           test_value=tmp_value.split('=')
           nbo_value = test_value[1] 
           thresh = float(nbo_value) 
       if 'eri_level' in lists[i]:
           tmp_value = str(lists[i])
           test_value=tmp_value.split('=')
           eri_level = int(test_value[1]) 
#   print('iop_nbo', iop_nbo)
   return iop_nbo_state, iop_nbo_act, iop_ri, iop_soc, thresh, eri_level
#=============================================================================
def cc_nbo(mol, dmab, mo_energy, mo_coeff,vrt, pop, isymmo, dirprd, thresh=0.0):
   frozen = 0
   iop_nbo_state = 0
   iop_nbo_act = 0
   iop_nbo_state, iop_nbo_act,_,_,thresh,_ = readinp()
   group_corr, nirrep, irrep_sym = symminfo(mol.groupname)

   nvrt = sum(vrt)
   nocc = sum(pop)
   orb_map = numpy.zeros(nocc+nvrt, dtype=int)
   irrep_map = numpy.zeros(nocc+nvrt, dtype=int)
   for irrep in range(nirrep):
       for i in range(pop[irrep]+vrt[irrep]):
           idx = isymmo[i, irrep]
           orb_map[idx] = i
           irrep_map[idx] = irrep

   ioffv = numpy.zeros(nirrep, dtype=int)
   dmvv_sym = numpy.zeros((nvrt,nvrt), dtype=float)
   irrep_n_sym = numpy.zeros(nvrt, dtype=float)
   n_sym = numpy.zeros(nvrt, dtype=float)
   v_sym = numpy.zeros((nvrt,nvrt), dtype=float)
#检查密度矩阵的对称性
   for i in range(nirrep-1):
       ioffv[i+1] = ioffv[i] + vrt[i]
   for sym_tmp in range(nirrep):
       for irrepb in range(nirrep):
           irrepa = dirprd[irrepb, sym_tmp]
           if irrepa!=irrepb:
               for b in range(vrt[irrepb]):
                   for a in range(vrt[irrepa]):
                       a_out = a + ioffv[irrepa]
                       b_out = b + ioffv[irrepb]
                       a_in = a + pop[irrepa]
                       b_in = b + pop[irrepb]
                       idxa = isymmo[a_in, irrepa]
                       idxb = isymmo[b_in, irrepb]
                       if dmab[idxa, idxb] > 1.e-10:
                           print('symerror in dmab', idxa, idxb, dmab[idxa, idxb])
           ioffva = ioffv[irrepa]
           ioffvb = ioffv[irrepb]
           numa = vrt[irrepa]
           numb = vrt[irrepb]
#按照对称性，重新排布密度矩阵
   for i in range(nirrep-1):
       ioffv[i+1] = ioffv[i] + vrt[i]
   for irrepb in range(nirrep):
       irrepa = dirprd[irrepb, 0]
       for b in range(vrt[irrepb]):
           for a in range(vrt[irrepa]):
               a_out = a + ioffv[irrepa]
               b_out = b + ioffv[irrepb]
               a_in = a + pop[irrepa]
               b_in = b + pop[irrepb]
               idxa = isymmo[a_in, irrepa]
               idxb = isymmo[b_in, irrepb]
               dmvv_sym[a_out,b_out] = dmab[idxa, idxb]
       ioffva = ioffv[irrepa]
       ioffvb = ioffv[irrepb]
       numa = vrt[irrepa]
       numb = vrt[irrepb]

#密度矩阵的对角化
       ntmp,vtmp= numpy.linalg.eigh(dmvv_sym[ioffva:ioffva+numa,ioffvb:ioffvb+numb])
       n_sym[ioffvb:ioffvb+numb]=ntmp
       v_sym[ioffva:ioffva+numa, ioffvb:ioffvb+numb]=vtmp
#验证密度矩阵的对称性
   for irrep in range(nirrep):
       if(irrep!=0):
           for irrepb in range(nirrep):
              irrepa = dirprd[irrepb, irrep]
              for b in range(vrt[irrepb]):
                  for a in range(vrt[irrepa]):
                      a_out = a + ioffv[irrepa]
                      b_out = b + ioffv[irrepb]
                      if(abs(dmvv_sym[a_out,b_out])>1.e-10):
                          print('error in symmetry of density V')
#                           exit()
#对不同的对称块的n和v按能量从高到低重新排布
   for irrep in range(nirrep):
       ntmp = n_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep]]
       vtmp = v_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep],ioffv[irrep]:ioffv[irrep]+vrt[irrep]]
       idx = numpy.argsort(ntmp)[::-1]
       ntmp,vtmp = ntmp[idx], vtmp[:,idx]
       irrep_n_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep]]                                = irrep
       n_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep]]                                      = ntmp
       v_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep],ioffv[irrep]:ioffv[irrep]+vrt[irrep]] = vtmp
   pct_occ = True

#   thresh = 5e-4
   idx = numpy.argsort(n_sym)[::1]
   ntmp = n_sym[idx]*1
   irrep_n_sym = irrep_n_sym[idx]
   for i in range(len(ntmp)):
       print('ncat check', i+nocc+1,'%16.8f'%ntmp[i], '%4i'%irrep_n_sym[i], 
                                    '%16.8f'%mo_energy[i+nocc], '%4i'%irrep_map[i+nocc])

# n_sym与原mo_energy相比，应优先满足简并度一致，再在此基础上由能量从低到高排序
   idx_mo = 0 
   igen_mo = 1 
   gennum_mo = numpy.zeros(nocc+nvrt, dtype=int)
   idx_nbo = 0 
   igen_nbo = 1 
   gennum_nbo = numpy.zeros(nocc+nvrt, dtype=int)
   length = len(mo_energy)
   for i in range(nocc,nocc+nvrt):
       if i<length-1:
           j = i - nocc
           if abs(ntmp[j+1]-ntmp[j])<1.e-10:
              igen_nbo = igen_nbo + 1
           else:
              gennum_nbo[idx_nbo]=igen_nbo
              idx_nbo = idx_nbo + 1  
              igen_nbo = 1 

           if abs(mo_energy[i+1]-mo_energy[i])<1.e-10:
              igen_mo = igen_mo + 1
           else:
              gennum_mo[idx_mo]=igen_mo
              idx_mo = idx_mo + 1  
              igen_mo = 1 
       else: 
           j = i - nocc
           if abs(mo_energy[i]-mo_energy[i-1])>1.e-10:
              igen_mo = 1 
              gennum_mo[idx_mo]=igen_mo
              idx_mo = idx_mo + 1
           if abs(ntmp[j]-ntmp[j-1])>1.e-10:
              igen_nbo = 1 
              gennum_nbo[idx]=igen_nbo
              idx_nbo = idx_nbo + 1
#   if idx_nbo!=idx_mo:
   print('attention!, error in degeneration', idx_nbo, idx_mo)
   genlth = idx_mo
   for i in range(genlth): 
       print('dgeneration', i, gennum_mo[i])  
   print('total numb', sum(gennum_mo))

   print('thresh', thresh, sum(n_sym))
   vrt_act = numpy.zeros(8,dtype = int)
#   if pct_occ is False:
   if iop_nbo_act==0:
#将占据数小于一定值的数组抛弃
      print('single act of nbo')
      for irrep in range(nirrep):
          ntmp = n_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep]]
          vrt_act[irrep] = numpy.count_nonzero(ntmp>thresh)
#将占据数总和小于一定值的数组抛弃
   elif iop_nbo_act==1:
      print('total  act of nbo')
      idx = numpy.argsort(n_sym)[::1]
      sumtmp = 0.0
      n_tmp = numpy.zeros(nvrt,dtype = int)
      n_tmp = n_sym[idx]
      for i in range(nvrt):
          sumtmp = sumtmp + n_tmp[i]
          if sumtmp<thresh:
             n_tmp[i] = 0.0
      n_sym[idx] = n_tmp
      for irrep in range(nirrep):
          ntmp = n_sym[ioffv[irrep]:ioffv[irrep]+vrt[irrep]]
          vrt_act[irrep] = numpy.count_nonzero(ntmp>0.0)

#将本征值重新排序

   #elif pct_occ is True:
#堆叠总比例大于一定数值
   print('finished vrt_act')
   ioffv2 = numpy.zeros(8, dtype=int)
   ioffv2[0] = 0
   for i in range(nirrep-1):
       ioffv2[i+1] = ioffv2[i] + vrt_act[i]
   nvrt_act = sum(vrt_act)
   v_act = numpy.zeros((nvrt,nvrt_act), dtype = float)
   for irrep in range(nirrep):
       for b in range(vrt_act[irrep]):
           b_in = b +  ioffv[irrep]
           b_out = b + ioffv2[irrep]
           v_act[:, b_out] = v_sym[:, b_in]

# 分子轨道能量按对称性重新排序    
   e_mo_sym = mo_energy % 10
   mo_coeff_sym =mo_coeff % 10
   idx = 0
   for irrep in range(nirrep):
       for i in range(pop[irrep]):
           iin = isymmo[i, irrep]
           e_mo_sym[idx] = mo_energy[iin]
           mo_coeff_sym[:,idx] = mo_coeff[:,iin]
           idx = idx + 1
   for irrep in range(nirrep):
       for i in range(vrt[irrep]):
           iin = isymmo[pop[irrep]+i, irrep]
           e_mo_sym[idx] = mo_energy[iin]
           mo_coeff_sym[:,idx] = mo_coeff[:,iin]
           idx = idx + 1
#
   fvv = numpy.diag(e_mo_sym[nocc:])
   fvv_no = numpy.dot(v_act.T, numpy.dot(fvv, v_act))
   n_canon = numpy.zeros(nvrt_act, dtype = float)
   v_canon = numpy.zeros((nvrt_act,nvrt_act), dtype = float)

   for irrepb in range(nirrep):
       irrepa = dirprd[irrepb, 0]
       ioffa = ioffv2[irrepa]
       numa = vrt_act[irrepa]
       n_canon_tmp, v_canon_tmp = numpy.linalg.eigh(fvv_no[ioffa:ioffa+numa,ioffa:ioffa+numa])
       n_canon[ioffa:ioffa+numa]=n_canon_tmp
       v_canon[ioffa:ioffa+numa, ioffa:ioffa+numa]=v_canon_tmp
#      print('n_tmp',n_canon_tmp) 
#      print('n_mo ',e_mo_sym[nocc+ioffa:nocc+ioffa+numa]) 
   print('finished v and n canon', nocc, nvrt, nvrt_act)
#   利用n_canon和v_canon去替换空轨道能量和空轨道系数
   no_coeff_1 = numpy.dot(mo_coeff_sym[:,nocc:], numpy.dot(v_act, v_canon))
   print('finished coeff1', no_coeff_1.shape)
#  no_coeff_2 = numpy.dot(mo_coeff_sym[:,nocc:], v_act)
#  print('finished coeff2', no_coeff_2.shape)
#  no_coeff_s = numpy.concatenate((mo_coeff_sym[:,:nocc], no_coeff_1, no_coeff_2), axis=1)
   no_coeff_s = numpy.concatenate((mo_coeff_sym[:,:nocc], no_coeff_1), axis=1)
   print('finished coeffs', no_coeff_s.shape)

#   escf = myhf.e_tot
   noa = sum(pop)
   nob = noa
   nva = sum(vrt_act)
   nvb = nva
   nmo = noa + nva

   idx = 0
   sym_id = numpy.zeros(noa+nva, dtype = int)
   for irrepi in range(nirrep):
       for i in range(pop[irrepi]):
           sym_id[idx] = irrepi
           idx = idx + 1
   for irrepi in range(nirrep):
       for i in range(pop[irrepi],pop[irrepi]+vrt_act[irrepi]):
           sym_id[idx] = irrepi
           idx = idx + 1

   e_mo_nbo = numpy.zeros(noa+nva, dtype = float)
   e_mo_nbo[:nocc]=e_mo_sym[:nocc]
   e_mo_nbo[nocc:nocc+nvrt_act] = n_canon[:nvrt_act]

   idx = numpy.argsort(e_mo_nbo)[::1]
   e_mo_nbo = e_mo_nbo[idx]
   no_coeff_s = no_coeff_s[:,idx]
   sym_id = sym_id[idx]
   e_mo_sym2 = numpy.sort(e_mo_sym)
   for i in range(noa+nva):
       print('e,symid',i,'%8.4f'%e_mo_sym2[i], '%8.4f'%e_mo_nbo[i], sym_id[i])

   orbn = no_coeff_s[:,frozen:]

   group_corr, nirrep, irrep_sym = symminfo(mol.groupname)
   symrx, symry, symrz = symrxyz(group_corr)

#   mo_sym.save_symm(dirprd, nirrep, nmo, vrt_act, pop, isymmo2, symrx, symry, symrz)
   return e_mo_nbo, no_coeff_s 

def symorb_sort_cf(mol, myhf, nirrep, mo_occ, mo_coeff, frozen=0, group_corr=0):
   listocc = list(numpy.argwhere(mo_occ > 0).flatten())
   listvir = list(numpy.argwhere(mo_occ == 0).flatten())
   nocc = len(listocc) - frozen
   nvir = len(listvir)
   nmo = nocc + nvir
   isympop = numpy.zeros((nmo,8), dtype = int)
   isymvrt = numpy.zeros((nmo,8), dtype = int)
   isymmo  = numpy.zeros((nmo,8), dtype = int)
   pop = numpy.zeros(8, dtype = int)
   vrt = numpy.zeros(8, dtype = int)
   if nirrep==1:
      pop[0]=nocc
      vrt[0]=nvir
      for i in range(0,nmo):
          isymmo[i,0]=i
   elif nirrep>1:
#      symname = symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mo_coeff)
      symid  = symm.label_orb_symm(mol, mol.irrep_id,    mol.symm_orb, mo_coeff)
      symid = symid %10
#   symname=numpy.zeros(nmo+frozen, dtype =(str,3))
#  symm_table = symm.param.IRREP_ID_TABLE[group_corr]
#  i = 0 
#  for key, value in symm_table.items():
#      symname[i] = key
#      i = i + 1

   lth = nocc+nvir+frozen
#  print('irrep name:', symname)
#  print('irrep   id:', symid)
#  print('irrep  id2:', symid2)
   nmo = nocc + nvir
   isymmo  = numpy.zeros((nmo,8), dtype = int)
   pop = numpy.zeros(8, dtype = int)
   vrt = numpy.zeros(8, dtype = int)

   symm_table = symm.param.IRREP_ID_TABLE[group_corr]
   length = len(symm_table)

   sym_name_py = numpy.zeros(length, dtype = (str,3))
   sym_id_py = numpy.zeros(length, dtype = int)
   sym_name_cf = numpy.zeros(length, dtype = (str,3))
   sym_id_cf = numpy.zeros(length, dtype = int)
   k = 0
   for key, value in symm_table.items():
#       print('key and value', key, value)
       sym_name_py[k] = key
       sym_id_py[k] = value
       k = k + 1


   print(nmo+frozen)
   symname_cfour=numpy.zeros(nmo+frozen, dtype =(str,3))
   symid_cfour=numpy.zeros(nmo+frozen, dtype =int)
   eig_cfour = numpy.zeros(nmo+frozen, dtype = float)
   print(symid_cfour)
#  with open('sym_cfour.data', mode='r') as fr:
#      outcfour=fr.readlines()
#       print('sym_cfour.data', len(outcfour)), len(symid_cfour)
#      for i in range(0, len(outcfour)):
#             print(outcfour[i])
#            symid_cfour[i]=outcfour[i][48]
#            symname_cfour[i]=outcfour[i][24:27]
#            eig_cfour[i] = outcfour[i][50:80]
#            print('%3i'%i,'%18.10f'%eig_cfour[i], symname_cfour[i], symid_cfour[i])
   with open('cfour.tmp', mode='r') as fr:
       outcfour=fr.readlines()
       print('cfour.tmp', len(outcfour))
       echeck=' MO #        E(hartree)    '
       ebegin='----   --------------------'
       for i in range(1, len(outcfour)):
           if str(ebegin) in str(outcfour[i]) and str(echeck) in str(outcfour[i-1]):
              print(str(outcfour[i]))
              k=0
              for j in range(i+1, i+nocc+frozen+1):
                  symid_cfour[k]=(outcfour[j][77])
                  symname_cfour[k]=(outcfour[j][72:75])
                  eig_cfour[k] = outcfour[j][19:35]
                  print('%3i'%(k+1),'%18.10f'%eig_cfour[k],
                        '%18.10f'%myhf.mo_energy[k],  symname_cfour[k], symid_cfour[k])
                  k=k+1
              print('-----------HOMO LUMO GAP------------------')
              for j in range(i+nocc+frozen+2, i+nocc+frozen+nvir+2):
#                 print('j', j, k, nocc+nvir+frozen)
                  symid_cfour[k]=(outcfour[j][77])
                  symname_cfour[k]=(outcfour[j][72:75])
                  eig_cfour[k] = outcfour[j][19:35]
                  print('%3i'%(k+1),'%18.10f'%eig_cfour[k],
                        '%18.10f'%myhf.mo_energy[k],  symname_cfour[k], symid_cfour[k])
                  k=k+1
              break

   symid_cfour = symid_cfour-1
   eig_pyscf_sym=numpy.zeros((nmo,8),dtype=float)
   eig_cfour_sym=numpy.zeros((nmo,8),dtype=float)
   lth_py=numpy.zeros(nirrep,dtype=int)
   lth_cf=numpy.zeros(nirrep,dtype=int)
   for j in range(nirrep):
       idx_py = 0
       idx_cf = 0
       for i in range(0, nmo+frozen):
         if symid[i] == j:
            eig_pyscf_sym[idx_py,j]=myhf.mo_energy[i]
            idx_py = idx_py + 1
         if symid_cfour[i] == j:
            eig_cfour_sym[idx_cf,j]=eig_cfour[i]
            sym_name_cf[j] = symname_cfour[i]
            sym_id_cf[j] = j
            idx_cf = idx_cf + 1
       lth_py[j] = idx_py
       lth_cf[j] = idx_cf
   for i in range(0, nirrep):
       sym_name_cf[i]=sym_name_cf[i].strip()
   print('=========================================================')
   print('====sym orbital number and sym_name of pyscf and cfour===')
   print('lth_py', lth_py)
   print('lth_cf', lth_cf)
   print('sym_name_py', sym_name_py)
   print(sym_id_py)
   print('sym_name_cf', sym_name_cf)
   print(sym_id_cf)

   iop_py = numpy.zeros(nirrep,dtype=int)
   iop_cf = numpy.zeros(nirrep,dtype=int)
   iop_select = 0
   iop_select_py = numpy.zeros(nirrep,dtype=int)
   iop_select_cf = numpy.zeros(nirrep,dtype=int)
   irrep_cfour = numpy.zeros(nirrep, dtype = int)
   irrep_cfour2 = numpy.zeros(nirrep, dtype = int)
   for irrep_py in range(nirrep):
       for irrep_cf in range(nirrep):
           iop_select=0
           if iop_select_py[irrep_py]==1 or iop_select_cf[irrep_cf]==1:
              continue
           if lth_py[irrep_py] != lth_cf[irrep_cf]:
              iop_select = 1 
           else: 
              for i in range(lth_py[irrep_py]):
                  c_py = eig_pyscf_sym[i,irrep_py]
                  c_cf = eig_cfour_sym[i,irrep_cf]
                  ttmp = c_py - c_cf 
                  ttmp = abs(ttmp/max(abs(c_py),abs(c_cf),0.00000000001)) 
#                  print('tmp', i, irrep_py, irrep_cf, c_py, c_cf,ttmp)
                  if ttmp>0.01:
                      iop_select = 1
#                      print('tmp', i, irrep_py, irrep_cf, c_py, c_cf,ttmp)
           
           if iop_select==0:
#               print('irrep pyscf and cfour', irrep_py, irrep_cf)
               iop_select_py[irrep_py] = 1
               iop_select_cf[irrep_cf] = 1
               irrep_cfour[irrep_py] = irrep_cf
   print('=========================================================')
   print('====to check the symmetry by orbital energy and irrep====')
   print('irrep_cfour1', irrep_cfour)
   print('irrep_namecf',sym_name_cf[irrep_cfour])
   print('irrep_lth_py',lth_py)
   print('irrep_lth_cf',lth_cf[irrep_cfour])
#  for irrep in range(nirrep): 
#     print('irrep', irrep, irrep_cfour[irrep])
#     print(eig_pyscf_sym[:lth_py[irrep],irrep])
#     print(eig_cfour_sym[:lth_py[irrep],irrep_cfour[irrep]])

   for irrep_py in range(nirrep):
       for irrep_cf in range(nirrep): 
           if sym_name_py[irrep_py] == sym_name_cf[irrep_cf]:
               irrep_cfour2[irrep_py] = irrep_cf
   iop_irrep = 0
   for irrep in range(nirrep):
       if lth_cf[irrep_cfour[irrep]]!=lth_cf[irrep_cfour2[irrep]]: 
           iop_irrep = 1
   if iop_irrep == 0:
       irrep_cfour=irrep_cfour2
       print('irrep_cfour_c',irrep_cfour2)
   print('=========================================================')
   print('====to check the symmetry by irrep name =================')
   print('irrep_cfour2',irrep_cfour2)
   print('irrep_namecf',sym_name_cf[irrep_cfour2])
   print('irrep_lth_py',lth_py)
   print('irrep_lth_cf',lth_cf[irrep_cfour2])
   print('=========================================================')
   print('attention! lth and energy criterion are more important   ')
   print('with equal lth and energy, irrepname is more important   ')
   print('=========================================================')
   for k in range(len(myhf.mo_energy)):
       print('%3i'%(k+1),'%18.10f'%eig_cfour[k], symid_cfour[k],
            '%18.10f'%myhf.mo_energy[k],  irrep_cfour[symid[k]])
   
   return irrep_cfour, sym_name_cf, symid, symid_cfour

def get_dirprd():
    dirprd=numpy.array(((0, 1, 2, 3, 4, 5, 6, 7),
             (1, 0, 3, 2, 5, 4, 7, 6),
             (2, 3, 0, 1, 6, 7, 4, 5),
             (3, 2, 1, 0, 7, 6, 5, 4),
             (4, 5, 6, 7, 0, 1, 2, 3),
             (5, 4, 7, 6, 1, 0, 3, 2),
             (6, 7, 4, 5, 2, 3, 0, 1),
             (7, 6, 5, 4, 3, 2, 1, 0)))
    return dirprd

def save_cc_pyscf_int(noa, nob, nva, nvb, nmo, escf, group_n, frozen, nirrep, irrep_sym, orb_energy_a, orb_energy_b):
    with open('cc_pyscf_int', 'w') as csvfile:
         file_cc = csv.writer(csvfile)
         file_cc.writerow([noa, nob, nva, nvb, nmo, escf, group_n,frozen])
         file_cc.writerow(orb_energy_a)
         file_cc.writerow(orb_energy_b)
         for i in range(nirrep):
            file_cc.writerow([irrep_sym[i], i])

def generate_eri(mol, myhf,pop,vrt,nf,nirrep,iop_ri, eri_level):
   coeff = myhf.mo_coeff
   orb = coeff[:,nf:]
   orbo = coeff[:,nf:nf+sum(pop)]
   if iop_ri==0 and eri_level == 0:
      mo = (orb, orb, orb, orb)
      ao2mo.outcore.general(mol, mo,  'ab_c', aosym='s2kl')

   elif iop_ri==0 and eri_level != 0:
      mo = (orbo , orb, orb, orb)
      ao2mo.outcore.general(mol, mo,  'ab_c', aosym='s2kl')

   elif iop_ri==1:
      naux_sym = ri_blocks.ri_blocks(myhf, nf, 'cd_b')
      norb = orb.shape[1]
      naux = numpy.sum(naux_sym[:nirrep])
      with open('dfpar.data', mode='w') as dfwrite:
           dfwrite.write('%8d%8d%8d\n' % (naux, norb, norb))
           for i in naux_sym:
               dfwrite.write('%8d\n' % i)

def generate_eri_uhf(mol, myhf,pop,vrt,nf,nirrep,iop_ri, eri_level):
   coeff_a = myhf.mo_coeff[0]
   orb_a = coeff_a[:,nf:]
   orbo_a = coeff_a[:,nf:nf+sum(pop[:,0])]

   coeff_b = myhf.mo_coeff[1]
   orb_b = coeff_b[:,nf:]
   orbo_b = coeff_b[:,nf:nf+sum(pop[:,1])]
   if iop_ri==0 and eri_level == 0:
      mo = (orb_a, orb_a, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'aa_c', aosym='s2kl')
      mo = (orb_b, orb_b, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'ab_c', aosym='s2kl')
      mo = (orb_b, orb_b, orb_b, orb_b)
      ao2mo.outcore.general(mol, mo,  'bb_c', aosym='s2kl')

   elif iop_ri==0 and eri_level != 0:
      mo = (orbo_a, orb_a, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'aa_c', aosym='s2kl')
      mo = (orbo_b, orb_b, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'ab_c', aosym='s2kl')
      mo = (orbo_b, orb_b, orb_b, orb_b)
      ao2mo.outcore.general(mol, mo,  'bb_c', aosym='s2kl')

   elif iop_ri==1:
      naux_sym = ri_blocks.uri_blocks(myhf, nf, 'cd_a', 'cd_b')
      norb_a = orb_a.shape[1]
      naux = numpy.sum(naux_sym[:nirrep])
      with open('dfpar.data', mode='w') as dfwrite:
           dfwrite.write('%8d%8d%8d\n' % (naux, norb_a, norb_a))
           for i in naux_sym:
               dfwrite.write('%8d\n' % i)

def generate_eri_rohf(mol, myhf,pop,vrt,nf,nirrep,iop_ri, eri_level):
#  print('myhf.mo_coeff a', myhf.mo_coeff)
#  print('myhf.mo_coeff b', myhf.mo_coeff)
   coeff_a = myhf.mo_coeff.ro[0]
   orb_a = coeff_a[:,nf:]
   orbo_a = coeff_a[:,nf:nf+sum(pop[:,0])]

   coeff_b = myhf.mo_coeff.ro[1]
   orb_b = coeff_b[:,nf:]
   orbo_b = coeff_b[:,nf:nf+sum(pop[:,1])]
   if iop_ri==0 and eri_level == 0:
      mo = (orb_a, orb_a, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'aa_c', aosym='s2kl')
      mo = (orb_b, orb_b, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'ab_c', aosym='s2kl')
      mo = (orb_b, orb_b, orb_b, orb_b)
      ao2mo.outcore.general(mol, mo,  'bb_c', aosym='s2kl')

   elif iop_ri==0 and eri_level != 0:
      mo = (orbo_a, orb_a, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'aa_c', aosym='s2kl')
      mo = (orbo_b, orb_b, orb_a, orb_a)
      ao2mo.outcore.general(mol, mo,  'ab_c', aosym='s2kl')
      mo = (orbo_b, orb_b, orb_b, orb_b)
      ao2mo.outcore.general(mol, mo,  'bb_c', aosym='s2kl')

   elif iop_ri==1:
      naux_sym = ri_blocks.rori_blocks(myhf, nf, 'cd_a', 'cd_b')
      norb_a = orb_a.shape[1]
      naux = numpy.sum(naux_sym[:nirrep])
      with open('dfpar.data', mode='w') as dfwrite:
           dfwrite.write('%8d%8d%8d\n' % (naux, norb_a, norb_a))
           for i in naux_sym:
               dfwrite.write('%8d\n' % i)


def start_nat_orb(mol, myhf, nirrep, thresh, nf):
   pop, vrt, isymmo = symorb_sort(mol, nirrep, myhf.mo_occ, myhf.mo_coeff,0)
   print('pop and vrt in nbo',sum(pop), sum(vrt))
   dmab = numpy.loadtxt('dmab.data')
   dirprd = get_dirprd()
   e_mo_nbo, mo_coeff_nbo = cc_nbo(mol, dmab, myhf.mo_energy,
               myhf.mo_coeff, vrt, pop, isymmo, dirprd, thresh)
   lth = len(e_mo_nbo)
   mo_occ_nbo = numpy.zeros(lth, dtype=int)
   mo_occ_nbo[0:sum(pop)] = 2

   mo_e_old = myhf.mo_energy

   myhf.mo_occ = mo_occ_nbo
   myhf.mo_coeff = mo_coeff_nbo
   myhf.mo_energy = e_mo_nbo
   pop_old = pop*1
   vrt_old = vrt*1
   isymmo_old = isymmo*1
   pop, vrt, isymmo = symorb_sort(mol, nirrep, myhf.mo_occ, myhf.mo_coeff,0)

   print('vrt and vrt_old', vrt, vrt_old)

   act_orb=numpy.zeros(sum(vrt), dtype=int)
   act_orb_irrep=numpy.zeros(sum(vrt), dtype=int)
   with open('acted_orbital', 'w') as actorb_file:
        k = 0
        for irrep in range(nirrep):
            lth = vrt_old[irrep] - vrt[irrep]
            if lth > 0 :
               for i in range(lth):
                   act_orb[k] = isymmo_old[pop[irrep]+vrt[irrep]+i,irrep]
                   act_orb_irrep[k] = irrep
                   k = k + 1
#                  print('acted orbital', isymmo_old[pop[irrep]+vrt[irrep]+i,irrep])
#                  actorb_file.write(str(isymmo_old[pop[irrep]+vrt[irrep]+i,irrep]))
        act_orb[0:k].sort()
        a_o2 = numpy.zeros(k, dtype=str)
        idxtmp = 0
        str_w=str()
        if(k>1):
           iop_tmp = 0
           iop_tmp2 =0
           for i in range(k-1):

               if act_orb[i+1] - act_orb[i] == 1:
                  iop_tmp = iop_tmp + 1
                  iop_tmp2 = 1
                  if iop_tmp==1:
                     str_tmp=str(act_orb[i]+1)+'>'
                  if i+1==k-1:
                     str_tmp = str_tmp + str(act_orb[i+1]+1)
                     print('str_tmp', str_tmp)
                     str_w = str_w + str_tmp
               elif act_orb[i+1] - act_orb[i] != 1:
                  if iop_tmp2 == 1:
                     str_tmp=str_tmp + str(act_orb[i]+1)
                  else:
                     str_tmp=str(act_orb[i]+1)
                  iop_tmp = 0
                  iop_tmp2 = 0
                  idxtmp = idxtmp+1
                  print('str_tmp',str_tmp)
                  str_w = str_w + str_tmp+'-'
        if nf>0:
           str_w='1>' + str(nf)+'-'+ str_w
        str_w = 'DROPMO=' + str_w
        print('str_w', str_w)
        actorb_file.write(str_w + '\n')
        pop, vrt, isymmo = symorb_sort(mol, nirrep, myhf.mo_occ, myhf.mo_coeff,nf)
   return vrt, pop, isymmo

def symorb_sort_ro(mol, nirrep, mo_occ, mo_coeff, frozen=0):
#   print(mo_occ)
#   print('mo_occ', mo_occ)
   occidx_a = mo_occ[frozen:] == 2
   occidx_a += mo_occ[frozen:] == 1
   viridx_a = mo_occ[frozen:] == 0
   occidx_b = mo_occ[frozen:] == 2
   viridx_b = mo_occ[frozen:] == 1 
   viridx_b += mo_occ[frozen:] == 0 
#  print('occ', occidx_a, occidx_b)
#  print('vrt', viridx_a, viridx_b)
   nmo = len(mo_occ) - frozen
   nocc_a = numpy.count_nonzero(occidx_a)
   nvir_a = numpy.count_nonzero(viridx_a)
   nocc_b = numpy.count_nonzero(occidx_b)
   nvir_b = numpy.count_nonzero(viridx_b)
#   isymmo_a  = numpy.zeros((nmo,8), dtype = int)
#   isymmo_b  = numpy.zeros((nmo,8), dtype = int)
   isymmo  = numpy.zeros((nmo,8,2), dtype = int)
   pop = numpy.zeros((8,2), dtype = int)
   vrt = numpy.zeros((8,2), dtype = int)

   if nirrep==1:
      pop[0,0]=nocc_a
      vrt[0,0]=nvir_a
      pop[0,1]=nocc_b
      vrt[0,1]=nvir_b
      isymmo[:,0,0] = numpy.arange(nmo)
      isymmo[:,0,1] = numpy.arange(nmo)
   elif nirrep>1:
#      print('mol.irrep_id', mol.irrep_id)
#     print('mol.symm_orb', mol.symm_orb[0])
#     print('mol.symm_orb', mol.symm_orb[1])
      symid_a = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mo_coeff[0,:,frozen:])
      symid_b = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mo_coeff[1,:,frozen:])
#      print(symid_a)
#      print(symid_b)
      for irrep in range(nirrep):
#=======define alpha spin
          idx0 = symid_a == irrep
          idx1 = occidx_a & idx0; idx2 = viridx_a & idx0
          irrep_occ = list(numpy.arange(nmo)[idx1])
          irrep_vir = list(numpy.arange(nmo)[idx2])
          irrep_mo  = irrep_occ + irrep_vir
          pop[irrep,0] = len(irrep_occ)
          vrt[irrep,0] = len(irrep_vir)
          isymmo[0:pop[irrep,0]+vrt[irrep,0], irrep,0] = irrep_mo
#          isymmo_a[0:pop[irrep,0]+vrt[irrep,0], irrep] = irrep_mo
#=======define beta spin
          idx0 = symid_b == irrep
          idx1 = occidx_b & idx0; idx2 = viridx_b & idx0
          irrep_occ = list(numpy.arange(nmo)[idx1])
          irrep_vir = list(numpy.arange(nmo)[idx2])
          irrep_mo  = irrep_occ + irrep_vir
          pop[irrep,1] = len(irrep_occ)
          vrt[irrep,1] = len(irrep_vir)
          isymmo[0:pop[irrep,0]+vrt[irrep,0], irrep,1] = irrep_mo
#          isymmo_b[0:pop[irrep,0]+vrt[irrep,0], irrep] = irrep_mo
#      print('alpha', isymmo_a)
#      print('beta ', isymmo_b)

   return pop, vrt, isymmo

